{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3XX46cTrh6iD"
      },
      "source": [
        "##### Copyright 2024 The AI Edge Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "sKrlWr6Kh-mF"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hST65kOHXpiL"
      },
      "source": [
        "# Retrain a speech recognition model with TensorFlow Lite Model Maker\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nShlCXGkbRVA"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/speech_recognition\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/speech_recognition.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/speech_recognition.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/speech_recognition.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BB5k6xNKJ5Xe"
      },
      "source": [
        "\n",
        "In this colab notebook, you'll learn how to use the [TensorFlow Lite Model Maker](https://www.tensorflow.org/lite/models/modify/model_maker) to train a speech recognition model that can classify spoken words or short phrases using one-second sound samples. The Model Maker library uses transfer learning to retrain an existing TensorFlow model with a new dataset, which reduces the amount of sample data and time required for training. \n",
        "\n",
        "By default, this notebook retrains the model (BrowserFft, from the [TFJS Speech Command Recognizer](https://github.com/tensorflow/tfjs-models/tree/master/speech-commands#speech-command-recognizer)) using a subset of words from the [speech commands dataset](https://www.tensorflow.org/datasets/catalog/speech_commands) (such as \"up,\" \"down,\" \"left,\" and \"right\"). Then it exports a TFLite model that you can run on a mobile device or embedded system (such as a Raspberry Pi). It also exports the trained model as a TensorFlow SavedModel.\n",
        "\n",
        "This notebook is also designed to accept a custom dataset of WAV files, uploaded to Colab in a ZIP file. The more samples you have for each class, the better your accuracy will be, but because the transfer learning process uses feature embeddings from the pre-trained model, you can still get a fairly accurate model with only a few dozen samples in each of your classes.\n",
        "\n",
        "**Note:** The model we'll be training is optimized for speech recognition with one-second samples. If you want to perform more generic audio classification (such as detecting different types of music), we suggest you instead follow [this Colab to retrain an audio classifier](https://colab.sandbox.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/audio_classification.ipynb).\n",
        "\n",
        "If you want to run the notebook with the default speech dataset, you can run the whole thing now by clicking **Runtime \u003e Run all** in the Colab toolbar. However, if you want to use your own dataset, then continue down to [Prepare the dataset](#scrollTo=cBsSAeYLkc1Z) and follow the instructions there.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AeZZ_cSsZfPx"
      },
      "source": [
        "### Import the required packages\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MelHQlE7FVue"
      },
      "source": [
        "You'll need TensorFlow, TFLite Model Maker, and some modules for audio manipulation, playback, and visualizations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wbMc4vHjaYdQ"
      },
      "outputs": [],
      "source": [
        "!sudo apt -y install libportaudio2\n",
        "!pip install tflite-model-maker"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rwUA9u4oWoCR"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import glob\n",
        "import random\n",
        "import shutil\n",
        "\n",
        "import librosa\n",
        "import soundfile as sf\n",
        "from IPython.display import Audio\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "import tensorflow as tf\n",
        "import tflite_model_maker as mm\n",
        "from tflite_model_maker import audio_classifier\n",
        "from tflite_model_maker.config import ExportFormat\n",
        "\n",
        "print(f\"TensorFlow Version: {tf.__version__}\")\n",
        "print(f\"Model Maker Version: {mm.__version__}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cBsSAeYLkc1Z"
      },
      "source": [
        "## Prepare the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NTTSXJxJq2Bz"
      },
      "source": [
        "To train with the default speech dataset, just run all the code below as-is.\n",
        "\n",
        "But if you want to train with your own speech dataset, follow these steps:\n",
        "\n",
        "**Note:** \n",
        "The model you'll retrain expects input data to be roughly one second of audio at 44.1 kHz. Model Maker perfoms automatic resampling for the training dataset, so there's no need to resample your dataset if it has a sample rate other than 44.1 kHz. But beware that audio samples longer than one second will be split into multiple one-second chunks, and the final chunk will be discarded if it's shorter than one second.\n",
        "\n",
        "1. Be sure each sample in your dataset is in **WAV file format, about one second long**. Then create a ZIP file with all your WAV files, organized into separate subfolders for each classification. For example, each sample for a speech command \"yes\" should be in a subfolder named \"yes\". Even if you have only one class, the samples must be saved in a subdirectory with the class name as the directory name. (This script assumes your dataset **is not split** into train/validation/test sets and performs that split for you.)\n",
        "2. Click the **Files** tab in the left panel and just drag-drop your ZIP file there to upload it.\n",
        "3. Use the following drop-down option to set **`use_custom_dataset`** to True.\n",
        "4. Then skip to [Prepare a custom audio dataset](#scrollTo=EobYerLQkiF1) to specify your ZIP filename and dataset directory name.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "AK9o98X7qyhU"
      },
      "outputs": [],
      "source": [
        "use_custom_dataset = False #@param [\"False\", \"True\"] {type:\"raw\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D2sNXbYVHjjy"
      },
      "source": [
        "### Generate a background noise dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gBVClNMwHtMD"
      },
      "source": [
        "Whether you're using the default speech dataset or a custom dataset, you should have a good set of background noises so your model can distinguish speech from other noises (including silence). \n",
        "\n",
        "Because the following background samples are provided in WAV files that are a minute long or longer, we need to split them up into smaller one-second samples so we can reserve some for our test dataset. We'll also combine a couple different sample sources to build a comprehensive set of background noises and silence:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qvJd9VfmHu29"
      },
      "outputs": [],
      "source": [
        "tf.keras.utils.get_file('speech_commands_v0.01.tar.gz',\n",
        "                        'http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz',\n",
        "                        cache_dir='./',\n",
        "                        cache_subdir='dataset-speech',\n",
        "                        extract=True)\n",
        "tf.keras.utils.get_file('background_audio.zip',\n",
        "                        'https://storage.googleapis.com/download.tensorflow.org/models/tflite/sound_classification/background_audio.zip',\n",
        "                        cache_dir='./',\n",
        "                        cache_subdir='dataset-background',\n",
        "                        extract=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7CAVFc3woB3_"
      },
      "source": [
        "**Note:** Although there is a newer version available, we're using v0.01 of the speech commands dataset because it's a smaller download. v0.01 includes 30 commands, while v0.02 adds five more (\"backward\", \"forward\", \"follow\", \"learn\", and \"visual\")."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xgwWNifGL-3b"
      },
      "outputs": [],
      "source": [
        "# Create a list of all the background wav files\n",
        "files = glob.glob(os.path.join('./dataset-speech/_background_noise_', '*.wav'))\n",
        "files = files + glob.glob(os.path.join('./dataset-background', '*.wav'))\n",
        "\n",
        "background_dir = './background'\n",
        "os.makedirs(background_dir, exist_ok=True)\n",
        "\n",
        "# Loop through all files and split each into several one-second wav files\n",
        "for file in files:\n",
        "  filename = os.path.basename(os.path.normpath(file))\n",
        "  print('Splitting', filename)\n",
        "  name = os.path.splitext(filename)[0]\n",
        "  rate = librosa.get_samplerate(file)\n",
        "  length = round(librosa.get_duration(filename=file))\n",
        "  for i in range(length - 1):\n",
        "    start = i * rate\n",
        "    stop = (i * rate) + rate\n",
        "    data, _ = sf.read(file, start=start, stop=stop)\n",
        "    sf.write(os.path.join(background_dir, name + str(i) + '.wav'), data, rate)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bVlvVq-SkeeO"
      },
      "source": [
        "### Prepare the speech commands dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q_q22T5UHbJG"
      },
      "source": [
        "We already downloaded the speech commands dataset, so now we just need to prune the number of classes for our model.\n",
        "\n",
        "This dataset includes over 30 speech command classifications, and most of them have over 2,000 samples. But because we're using transfer learning, we don't need that many samples. So the following code does a few things:\n",
        "\n",
        "+ Specify which classifications we want to use, and delete the rest.\n",
        "+ Keep only 150 samples of each class for training (to prove that transfer learning works well with smaller datasets and simply to reduce the training time).\n",
        "+ Create a separate directory for a test dataset so we can easily run inference with them later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zUSRpw2nOp8p"
      },
      "outputs": [],
      "source": [
        "if not use_custom_dataset:\n",
        "  commands = [ \"up\", \"down\", \"left\", \"right\", \"go\", \"stop\", \"on\", \"off\", \"background\"]\n",
        "  dataset_dir = './dataset-speech'\n",
        "  test_dir = './dataset-test'\n",
        "\n",
        "  # Move the processed background samples\n",
        "  shutil.move(background_dir, os.path.join(dataset_dir, 'background'))   \n",
        "\n",
        "  # Delete all directories that are not in our commands list\n",
        "  dirs = glob.glob(os.path.join(dataset_dir, '*/'))\n",
        "  for dir in dirs:\n",
        "    name = os.path.basename(os.path.normpath(dir))\n",
        "    if name not in commands:\n",
        "      shutil.rmtree(dir)\n",
        "\n",
        "  # Count is per class\n",
        "  sample_count = 150\n",
        "  test_data_ratio = 0.2\n",
        "  test_count = round(sample_count * test_data_ratio)\n",
        "\n",
        "  # Loop through child directories (each class of wav files)\n",
        "  dirs = glob.glob(os.path.join(dataset_dir, '*/'))\n",
        "  for dir in dirs:\n",
        "    files = glob.glob(os.path.join(dir, '*.wav'))\n",
        "    random.seed(42)\n",
        "    random.shuffle(files)\n",
        "    # Move test samples:\n",
        "    for file in files[sample_count:sample_count + test_count]:\n",
        "      class_dir = os.path.basename(os.path.normpath(dir))\n",
        "      os.makedirs(os.path.join(test_dir, class_dir), exist_ok=True)\n",
        "      os.rename(file, os.path.join(test_dir, class_dir, os.path.basename(file)))\n",
        "    # Delete remaining samples\n",
        "    for file in files[sample_count + test_count:]:\n",
        "      os.remove(file)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EobYerLQkiF1"
      },
      "source": [
        "### Prepare a custom dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3xTvDP3knMd"
      },
      "source": [
        "If you want to train the model with our own speech dataset, you need to upload your samples as WAV files in a ZIP ([as described above](#scrollTo=cBsSAeYLkc1Z)) and modify the following variables to specify your dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77PsQAKA4Arx"
      },
      "outputs": [],
      "source": [
        "if use_custom_dataset:\n",
        "  # Specify the ZIP file you uploaded:\n",
        "  !unzip YOUR-FILENAME.zip\n",
        "  # Specify the unzipped path to your custom dataset\n",
        "  # (this path contains all the subfolders with classification names):\n",
        "  dataset_dir = './YOUR-DIRNAME'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hwp6EQqvttgf"
      },
      "source": [
        "After changing the filename and path name above, you're ready to train the model with your custom dataset. In the Colab toolbar, select **Runtime \u003e Run all** to run the whole notebook.\n",
        "\n",
        "The following code integrates our new background noise samples into your dataset and then separates a portion of all samples to create a test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tMQ6cpw_B9e_"
      },
      "outputs": [],
      "source": [
        "def move_background_dataset(dataset_dir):\n",
        "  dest_dir = os.path.join(dataset_dir, 'background')\n",
        "  if os.path.exists(dest_dir):\n",
        "    files = glob.glob(os.path.join(background_dir, '*.wav'))\n",
        "    for file in files:\n",
        "      shutil.move(file, dest_dir)\n",
        "  else:\n",
        "    shutil.move(background_dir, dest_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "45iru8OdliG3"
      },
      "outputs": [],
      "source": [
        "if use_custom_dataset:\n",
        "  # Move background samples into custom dataset\n",
        "  move_background_dataset(dataset_dir)\n",
        "\n",
        "  # Now we separate some of the files that we'll use for testing:\n",
        "  test_dir = './dataset-test'\n",
        "  test_data_ratio = 0.2\n",
        "  dirs = glob.glob(os.path.join(dataset_dir, '*/'))\n",
        "  for dir in dirs:\n",
        "    files = glob.glob(os.path.join(dir, '*.wav'))\n",
        "    test_count = round(len(files) * test_data_ratio)\n",
        "    random.seed(42)\n",
        "    random.shuffle(files)\n",
        "    # Move test samples:\n",
        "    for file in files[:test_count]:\n",
        "      class_dir = os.path.basename(os.path.normpath(dir))\n",
        "      os.makedirs(os.path.join(test_dir, class_dir), exist_ok=True)\n",
        "      os.rename(file, os.path.join(test_dir, class_dir, os.path.basename(file)))\n",
        "    print('Moved', test_count, 'images from', class_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "myPa1dfEoagz"
      },
      "source": [
        "### Play a sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jecBYREgMk6"
      },
      "source": [
        "To be sure the dataset looks correct, let's play at a random sample from the test set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gLC3ayJsoeNw"
      },
      "outputs": [],
      "source": [
        "def get_random_audio_file(samples_dir):\n",
        "  files = os.path.abspath(os.path.join(samples_dir, '*/*.wav'))\n",
        "  files_list = glob.glob(files)\n",
        "  random_audio_path = random.choice(files_list)\n",
        "  return random_audio_path\n",
        "\n",
        "def show_sample(audio_path):\n",
        "  audio_data, sample_rate = sf.read(audio_path)\n",
        "  class_name = os.path.basename(os.path.dirname(audio_path))\n",
        "  print(f'Class: {class_name}')\n",
        "  print(f'File: {audio_path}')\n",
        "  print(f'Sample rate: {sample_rate}')\n",
        "  print(f'Sample length: {len(audio_data)}')\n",
        "\n",
        "  plt.title(class_name)\n",
        "  plt.plot(audio_data)\n",
        "  display(Audio(audio_data, rate=sample_rate))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "todbtEWFy0mj"
      },
      "outputs": [],
      "source": [
        "random_audio = get_random_audio_file(test_dir)\n",
        "show_sample(random_audio)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f-jRIWcQv7xt"
      },
      "source": [
        "## Define the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pQj1Mf7YZELS"
      },
      "source": [
        "\n",
        "When using Model Maker to retrain any model, you have to start by defining a model spec. The spec defines the base model from which your new model will extract feature embeddings to begin learning new classes. The spec for this speech recognizer is based on the pre-trained [BrowserFft model from TFJS](https://github.com/tensorflow/tfjs-models/tree/master/speech-commands#speech-command-recognizer).\n",
        "\n",
        "The model expects input as an audio sample that's 44.1 kHz, and just under a second long: the exact sample length must be 44034 frames.\n",
        "\n",
        "You don't need to do any resampling with your training dataset. Model Maker takes care of that for you. But when you later run inference, you must be sure that your input matches that expected format.\n",
        "\n",
        "All you need to do here is instantiate the [`BrowserFftSpec`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/BrowserFftSpec):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tUcxtfHXY7XS"
      },
      "outputs": [],
      "source": [
        "spec = audio_classifier.BrowserFftSpec()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "maTOoRvAwI9l"
      },
      "source": [
        "## Load your dataset "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UASCEHoVwQ1q"
      },
      "source": [
        "Now you need to load your dataset according to the model specifications. Model Maker includes the [`DataLoader`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/DataLoader) API, which will load your dataset from a folder and ensure it's in the expected format for the model spec.\n",
        "\n",
        "We already reserved some test files by moving them to a separate directory, which makes it easier to run inference with them later. Now we'll create a `DataLoader` for each split: the training set, the validation set, and the test set."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bAhAfHwiw2_F"
      },
      "source": [
        "#### Load the speech commands dataset\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cX0RqETqZgzo"
      },
      "outputs": [],
      "source": [
        "if not use_custom_dataset:\n",
        "  train_data_ratio = 0.8\n",
        "  train_data = audio_classifier.DataLoader.from_folder(\n",
        "      spec, dataset_dir, cache=True)\n",
        "  train_data, validation_data = train_data.split(train_data_ratio)\n",
        "  test_data = audio_classifier.DataLoader.from_folder(\n",
        "      spec, test_dir, cache=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2OWQ_O9_t-C-"
      },
      "source": [
        "#### Load a custom dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IPiwaJwMt7yo"
      },
      "source": [
        "**Note:** Setting `cache=True` is important to make training faster (especially when the dataset must be re-sampled) but it will also require more RAM to hold the data. If you use a very large custom dataset, caching might exceed your RAM capacity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "code",
        "id": "e86Ej-ZmuCzy"
      },
      "outputs": [],
      "source": [
        "if use_custom_dataset:\n",
        "  train_data_ratio = 0.8\n",
        "  train_data = audio_classifier.DataLoader.from_folder(\n",
        "      spec, dataset_dir, cache=True)\n",
        "  train_data, validation_data = train_data.split(train_data_ratio)\n",
        "  test_data = audio_classifier.DataLoader.from_folder(\n",
        "      spec, test_dir, cache=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hh1P_zfzwbfE"
      },
      "source": [
        "## Train the model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ziMghju-Rts2"
      },
      "source": [
        "Now we'll use the Model Maker [`create()`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/create) function to create a model based on our model spec and training dataset, and begin training.\n",
        "\n",
        "If you're using a custom dataset, you might want to change the batch size as appropriate for the number of samples in your train set.\n",
        "\n",
        "**Note:** The first epoch takes longer because it must create the cache. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GYaZvaOPgLUC"
      },
      "outputs": [],
      "source": [
        "# If your dataset has fewer than 100 samples per class,\n",
        "# you might want to try a smaller batch size\n",
        "batch_size = 25\n",
        "epochs = 25\n",
        "model = audio_classifier.create(train_data, spec, validation_data, batch_size, epochs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mtLuRA2xweZA"
      },
      "source": [
        "## Review the model performance"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oXMEHZkAxJTl"
      },
      "source": [
        "Even if the accuracy/loss looks good from the training output above, it's important to also run the model using test data that the model has not seen yet, which is what the `evaluate()` method does here:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n_4MGpzhWVhr"
      },
      "outputs": [],
      "source": [
        "model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HqvpAnqsVExO"
      },
      "source": [
        "### View the confusion matrix"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8QRRAM39aOxS"
      },
      "source": [
        "When training a classification model such as this one, it's also useful to inspect the [confusion matrix](https://en.wikipedia.org/wiki/Confusion_matrix). The confusion matrix gives you detailed visual representation of how well your classifier performs for each classification in your test data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zqB3c0368iH3"
      },
      "outputs": [],
      "source": [
        "def show_confusion_matrix(confusion, test_labels):\n",
        "  \"\"\"Compute confusion matrix and normalize.\"\"\"\n",
        "  confusion_normalized = confusion.astype(\"float\") / confusion.sum(axis=1)\n",
        "  sns.set(rc = {'figure.figsize':(6,6)})\n",
        "  sns.heatmap(\n",
        "      confusion_normalized, xticklabels=test_labels, yticklabels=test_labels,\n",
        "      cmap='Blues', annot=True, fmt='.2f', square=True, cbar=False)\n",
        "  plt.title(\"Confusion matrix\")\n",
        "  plt.ylabel(\"True label\")\n",
        "  plt.xlabel(\"Predicted label\")\n",
        "\n",
        "confusion_matrix = model.confusion_matrix(test_data)\n",
        "show_confusion_matrix(confusion_matrix.numpy(), test_data.index_to_label)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yASrikBgZ9ZO"
      },
      "source": [
        "## Export the model\n",
        "\n",
        "The last step is exporting your model into the TensorFlow Lite format for execution on mobile/embedded devices and into the [SavedModel format](https://www.tensorflow.org/guide/saved_model) for execution elsewhere.\n",
        "\n",
        "When exporting a `.tflite` file from Model Maker, it includes [model metadata](https://www.tensorflow.org/lite/inference_with_metadata/overview) that describes various details that can later help during inference. It even includes a copy of the classification labels file, so you don't need to a separate `labels.txt` file. (In the next section, we show how to use this metadata to run an inference.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4gEf59NfGWjq"
      },
      "outputs": [],
      "source": [
        "TFLITE_FILENAME = 'browserfft-speech.tflite'\n",
        "SAVE_PATH = './models'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xw_ehPxAdQlz"
      },
      "outputs": [],
      "source": [
        "print(f'Exporing the model to {SAVE_PATH}')\n",
        "model.export(SAVE_PATH, tflite_filename=TFLITE_FILENAME)\n",
        "model.export(SAVE_PATH, export_format=[mm.ExportFormat.SAVED_MODEL, mm.ExportFormat.LABEL])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lIC1ddGq6xQX"
      },
      "source": [
        "## Run inference with TF Lite model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5xr0idac6xfi"
      },
      "source": [
        "Now your TFLite model can be deployed and run using any of the supported [inferencing libraries](https://www.tensorflow.org/lite/guide/inference) or with the new [TFLite AudioClassifier Task API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/audio_classifier). The following code shows how you can run inference with the `.tflite` model in Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nR5zV53YbCIQ"
      },
      "outputs": [],
      "source": [
        "# This library provides the TFLite metadata API\n",
        "! pip install -q tflite_support"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1AC7PRyiayU5"
      },
      "outputs": [],
      "source": [
        "from tflite_support import metadata\n",
        "import json\n",
        "\n",
        "def get_labels(model):\n",
        "  \"\"\"Returns a list of labels, extracted from the model metadata.\"\"\"\n",
        "  displayer = metadata.MetadataDisplayer.with_model_file(model)\n",
        "  labels_file = displayer.get_packed_associated_file_list()[0]\n",
        "  labels = displayer.get_associated_file_buffer(labels_file).decode()\n",
        "  return [line for line in labels.split('\\n')]\n",
        "\n",
        "def get_input_sample_rate(model):\n",
        "  \"\"\"Returns the model's expected sample rate, from the model metadata.\"\"\"\n",
        "  displayer = metadata.MetadataDisplayer.with_model_file(model)\n",
        "  metadata_json = json.loads(displayer.get_metadata_json())\n",
        "  input_tensor_metadata = metadata_json['subgraph_metadata'][0][\n",
        "          'input_tensor_metadata'][0]\n",
        "  input_content_props = input_tensor_metadata['content']['content_properties']\n",
        "  return input_content_props['sample_rate']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yC7TEvQ9o4mu"
      },
      "source": [
        "To observe how well the model performs with real samples, run the following code block over and over. Each time, it will fetch a new test sample and run inference with it, and you can listen to the audio sample below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "loU6PleipSPf"
      },
      "outputs": [],
      "source": [
        "# Get a WAV file for inference and list of labels from the model\n",
        "tflite_file = os.path.join(SAVE_PATH, TFLITE_FILENAME)\n",
        "labels = get_labels(tflite_file)\n",
        "random_audio = get_random_audio_file(test_dir)\n",
        "\n",
        "# Ensure the audio sample fits the model input\n",
        "interpreter = tf.lite.Interpreter(tflite_file)\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "input_size = input_details[0]['shape'][1]\n",
        "sample_rate = get_input_sample_rate(tflite_file)\n",
        "audio_data, _ = librosa.load(random_audio, sr=sample_rate)\n",
        "if len(audio_data) \u003c input_size:\n",
        "  audio_data.resize(input_size)\n",
        "audio_data = np.expand_dims(audio_data[:input_size], axis=0)\n",
        "\n",
        "# Run inference\n",
        "interpreter.allocate_tensors()\n",
        "interpreter.set_tensor(input_details[0]['index'], audio_data)\n",
        "interpreter.invoke()\n",
        "output_data = interpreter.get_tensor(output_details[0]['index'])\n",
        "\n",
        "# Display prediction and ground truth\n",
        "top_index = np.argmax(output_data[0])\n",
        "label = labels[top_index]\n",
        "score = output_data[0][top_index]\n",
        "print('---prediction---')\n",
        "print(f'Class: {label}\\nScore: {score}')\n",
        "print('----truth----')\n",
        "show_sample(random_audio)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VtmfoJW6G2fd"
      },
      "source": [
        "## Download the TF Lite model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8zLDeiQ_z2Vj"
      },
      "source": [
        "Now you can deploy the TF Lite model to your mobile or embedded device. You don't need to download the labels file because you can instead retrieve the labels from `.tflite` file metadata, as shown in the previous inferencing example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cNuQoqtjG4zu"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  from google.colab import files\n",
        "except ImportError:\n",
        "  pass\n",
        "else:\n",
        "  files.download(tflite_file)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iERuGZz4z6rB"
      },
      "source": [
        "Check out our end-to-end example apps that perform inferencing with TFLite audio models on [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/sound_classification/android/) and [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/sound_classification/ios)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "model_maker_speech_recognition.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
